from pyted import *
from pyted.modules import *
from lut.lookuptable.MOSinfo import LUT_DB
from tedhub.common.opamp import OpampGenerator
from tedhub.circuit.testers import OTATester, RingAmpTester, TIATester
from taco.ea import DE


def gmid_LUT(Vds, Vbs, ID, model, gmid, gm_gds_req, l):
    if model == "nch":
        reverse_v = -1
    else:
        reverse_v = 1
    flag = 0
    """calculate L and W for M1"""
    for i in l:
        vgs = LUT_DB().lookup_Vgs(
            model,
            L=i,
            VDS=Vds,
            VBS=Vbs*reverse_v,
            GMOVERID=gmid,)

        # PMOS 的 VGS 和 VDS 应为负数， VBS 为正数
        gm_gds = LUT_DB().lookup(
            model=model,
            outpar="SELF_GAIN",
            VDS=-Vds*reverse_v,
            VGS=-vgs*reverse_v,
            VBS=Vbs*reverse_v,
            L=i,)[0]
        # debug(vgs, gmid)
        if gm_gds > gm_gds_req:
            # 输出满足条件的第一个栅长值，也就是最小栅长
            flag = 1
            L1 = round(i, 4)
            Vgs1 = vgs
            # debug(f"选择的栅长为:L={L1}")
            break

    if flag == 0:
        debug("未找到满足条件的栅长!")

    # 计算栅宽
    jd1 = (
        LUT_DB().lookup(
            model=model,
            outpar="ID",
            VDS=-Vds*reverse_v,
            VGS=-Vgs1*reverse_v,
            VBS=Vbs*reverse_v,
            L=L1,
        )[0]
        / 5
    )

    W1 = round(ID / jd1, 2)
    gmid1 = LUT_DB().lookup(
        model=model,
        outpar="GMOVERID",
        VDS=-Vds*reverse_v,
        VGS=-Vgs1*reverse_v,
        VBS=0,
        L=L1,)
    # 至此 M1 的 W 和 L 设计完毕
    # debug(f"待查MOS的栅宽为:{W1},栅长为:{L1},gmid 为:{gmid}")
    return [W1, L1, gmid1, gm_gds, ID]


def w_regulated(w_input, w_set):
    if w_input < w_set:
        w_output = w_input
        multi = 1
    else:
        multi = int[np.ceil(w_input / w_set)]
        w_output = w_input / multi
    return [w_output, multi]


class FDAmp(OpampGenerator):
    def __init__(self, vdd, Av, GBW, SR, PM_deg, power, Cl, **kwargs):
        super().__init__("dido")
        self.vdd = vdd  # 1.2
        self.Av = Av  # 10000
        self.GBW = GBW  # 100e6
        self.SR = SR  # 30
        self.PM_deg = PM_deg  # 63
        self.power = power  # 2e-3
        self.cl = Cl  # 3e-12

        self.para_Cc = 1.9  # #微调使用    t65:1.9; Except t65: 2.1
        self.w_set = 20  # 版图使用

        VDD = self.vdd
        Power = self.power
        Av = self.Av
        PM_deg = self.PM_deg
        para_Cc = self.para_Cc
        CL = self.cl
        w_set = self.w_set

        PM_rad = PM_deg*np.pi/180
        PM = np.tan(PM_rad)
        Vds = VDD/5
        ID = Power/63/VDD
        gmgds1 = 0  # pow(Av/15*2, 0.5)
        debug(f"PM:{PM}, gmgds1:{gmgds1}")
        debug(f"Vds:{Vds}, ID:{ID}")

        gmid1 = 20
        gmid2 = 15
        cc = 2*gmid1/gmid2/6*CL*para_Cc
        rc = 1/gmid2/ID/24*(1+3e-12/cc)

        l = np.arange(0.2, 2, 0.1)

        Mp = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
                      gmid=15, gm_gds_req=20, l=l)
        Mp1 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
                       gmid=gmid1, gm_gds_req=gmgds1, l=l)
        Mp2 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
                       gmid=18, gm_gds_req=gmgds1, l=l)
        Mn1 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="nch",
                       gmid=18, gm_gds_req=gmgds1, l=l)
        Mn = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="nch",
                      gmid=gmid2, gm_gds_req=gmgds1, l=l)
        debug(f"Mp  w={Mp[0]} ,l={Mp[1]} ,gmid={Mp[2]}")
        debug(f"Mp1 w={Mp1[0]},l={Mp1[1]},gmid={Mp1[2]}")
        debug(f"Mp2 w={Mp2[0]},l={Mp2[1]},gmid={Mp2[2]}")
        debug(f"Mn1 w={Mn1[0]},l={Mn1[1]},gmid={Mn1[2]}")
        debug(f"Mn  w={Mn[0]} ,l={Mn[1]} ,gmid={Mn[2]}")
        debug(f"Cc={cc}, Rm={rc}")

        u = 1e-6
        self.R_and_C = (
            Variable(name="cc", value=cc*1e12).set_bound(
                0.3*cc*1e12, 3*cc*1e12)*1e-12,
            Variable(name="rc", value=rc*1e-2).set_bound(
                0.3*rc*1e-2, 3*rc*1e-2)*1e2,
        )

        self.Mn_opt = (
            Variable(name="wn", value=Mn[0]).set_bound(
                max(0.3*Mn[0], self.wmin*1e6), min(3*Mn[0], self.wmax*1e6)) * u,
            Variable(name="ln", value=Mn[1]).set_bound(
                0.2, min(3*Mn[1], self.lmax*1e6)) * u,
        )

        self.Mn1_opt = (
            Variable(name="wn1", value=Mn1[0]).set_bound(
                max(0.3*Mn1[0], self.wmin*1e6), min(3*Mn1[0], self.wmax*1e6)) * u,
            Variable(name="ln1", value=Mn1[1]).set_bound(
                0.2, min(3*Mn1[1], self.lmax*1e6)) * u,
        )

        self.Mp_opt = (
            Variable(name="wp", value=Mp[0]).set_bound(
                max(0.3*Mp[0], self.wmin*1e6), min(3*Mp[0], self.wmax*1e6)) * u,
            Variable(name="lp", value=Mp[1]).set_bound(
                0.2, min(3*Mp[1], self.lmax*1e6)) * u,
        )

        self.Mp1_opt = (
            Variable(name="wp1", value=Mp1[0]).set_bound(
                max(0.3*Mp1[0], self.wmin*1e6), min(3*Mp1[0], self.wmax*1e6)) * u,
            Variable(name="lp1", value=Mp1[1]).set_bound(
                0.2, min(3*Mp1[1], self.lmax*1e6)) * u,
        )

        self.Mp2_opt = (
            Variable(name="wp2", value=Mp2[0]).set_bound(
                max(0.3*Mp2[0], self.wmin*1e6), min(3*Mp2[0], self.wmax*1e6)) * u,
            Variable(name="lp2", value=Mp2[1]).set_bound(
                0.2, min(3*Mp2[1], self.lmax*1e6)) * u,
        )

    def core(
        self, Vip: Input, Vin: Input, Idc: Ctrl, Vcm: Ctrl, Vop: Output, Von: Output
    ):
        Vm1, Vm2, Vm4, Vm5, Vm6, Vm7, Vm8, VB2 = Node(8)
        Vtail, Vf1, Vf2, Vcmfb = Node(4)

        """性能指标"""
        VDD = self.vdd
        Power = self.power
        Av = self.Av
        PM_deg = self.PM_deg
        para_Cc = self.para_Cc
        CL = self.cl
        w_set = self.w_set

        PM_rad = PM_deg*np.pi/180
        PM = np.tan(PM_rad)
        Vds = VDD/5
        ID = Power/63/VDD
        gmgds1 = pow(Av/15*2, 0.5)
        debug(f"PM:{PM}, gmgds1:{gmgds1}")
        debug(f"Vds:{Vds}, ID:{ID}")

        gmid1 = 20
        gmid2 = 15
        # cc = 2*gmid1/gmid2/6*CL*para_Cc
        # rc = 1/gmid2/ID/24*(1+3e-12/cc)

        # l = np.arange(0.2, 2, 0.1)

        # Mp = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
        #               gmid=15, gm_gds_req=20, l=l)
        # Mp1 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
        #                gmid=gmid1, gm_gds_req=gmgds1, l=l)
        # Mp2 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="pch",
        #                gmid=18, gm_gds_req=gmgds1, l=l)
        # Mn1 = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="nch",
        #                gmid=18, gm_gds_req=gmgds1, l=l)
        # Mn = gmid_LUT(Vds=Vds, Vbs=0, ID=ID, model="nch",
        #               gmid=gmid2, gm_gds_req=gmgds1, l=l)
        # debug(f"Mp  w={Mp[0]} ,l={Mp[1]} ,gmid={Mp[2]}")
        # debug(f"Mp1 w={Mp1[0]},l={Mp1[1]},gmid={Mp1[2]}")
        # debug(f"Mp2 w={Mp2[0]},l={Mp2[1]},gmid={Mp2[2]}")
        # debug(f"Mn1 w={Mn1[0]},l={Mn1[1]},gmid={Mn1[2]}")
        # debug(f"Mn  w={Mn[0]} ,l={Mn[1]} ,gmid={Mn[2]}")
        # debug(f"Cc={cc}, Rm={rc}")
        cc, rc = self.R_and_C
        wn, ln = self.Mn_opt
        wn1, ln1 = self.Mn1_opt
        wp, lp = self.Mp_opt
        wp1, lp1 = self.Mp1_opt
        wp2, lp2 = self.Mp2_opt
        [wn, multi_n] = w_regulated(wn, w_set)
        [wn1, multi_n1] = w_regulated(wn1, w_set)
        [wp, multi_p] = w_regulated(wp, w_set)
        [wp1, multi_p1] = w_regulated(wp1, w_set)
        [wp2, multi_p2] = w_regulated(wp2, w_set)
        nf = 2
        # [wn, multi_n] = [wn/nf, multi_n * nf]
        [wn1, multi_n1] = [wn1/nf, multi_n1 * nf]
        # [wp, multi_p] = [wp/nf, multi_p*nf]
        [wp1, multi_p1] = [wp1/nf, multi_p1*nf]
        [wp2, multi_p2] = [wp2/nf, multi_p2*nf]
        # debug(wn, ln)

        cmn1 = (
            CurrentMirror(l=ln, w=wn, m=1*multi_n, model="nch", type="type1", place_mode="euler")
            .input(Idc)
            .output(Vm1)
            .output(VB2)
        )
        cmn2b = CurrentMirror(l=[ln1, ln * 3], w=[wn1, wn], m=[multi_n1, multi_n], 
                              model="nch", type="type4", place_mode="euler").input(Vm4)
        cmn2 = CurrentMirror(l=[ln1, ln], w=[wn1, wn],
                             model="nch", type="type5", place_mode="euler").input(Vm5, Vm4)

        cmp1b = CurrentMirror(l=[lp1, lp * 3], w=[wp1, wp], m=[multi_p1, multi_p], 
                              model="pch", type="type4", place_mode="euler").input(Vm1)
        cmp1 = (
            CurrentMirror(l=[lp1, lp], w=[wp1, wp], m=[multi_p1, multi_p], model="pch", type="type5", place_mode="euler")
            .input(Vm2, Vm1)
            .output(Vm4)
        )

        "first stage"
        # first stage
        (
            (Vtail, Vtail)
            >> (dp1 := ABBA(l=lp1, w=wp1, m=2*multi_p1, model="pch")) % (Vip, Vin)
            >> (Vf1, Vf2)
            >> (dp2 := ABBA(l=lp2, w=wp2, m=3*multi_p2, model="pch")) % (VB2, VB2)
            >> (Vm7, Vm8)
            >> (dp3 := ABBA(l=ln1, w=wn1, m=3*multi_n1, model="nch")) % (Vm4, Vm4)
            >> (dp4 := ABBA(l=ln, w=wn, m=3*multi_n, model="nch")) % (Vm5, Vm5)
            >> (Vss(), Vss())
        )

        # second stage
        (
            (Vdd(), Vdd())
            >> (dp5 := ABBA(l=lp, w=wp, m=12*multi_p, model="pch")) % (Vm2, Vm2)
            >> (Von, Vop)
            >> (dp6 := ABBA(l=ln, w=wn, m=12*multi_n, model="nch")) % (Vm8, Vm7)
            >> (Vss(), Vss())
        )

        # cmfb
        (Vm6, Vm6) >> (dp7 := ABBA(l=lp1, w=wp1, m=1*multi_p1, model="pch") %
                       (Vcm, Vcm)) >> (Vf1, Vf2)
        R(Vop, Vcmfb, r=100e3)
        R(Von, Vcmfb, r=100e3)
        C(Vop, Vcmfb, c=100e-15)
        C(Von, Vcmfb, c=100e-15)

        R(Vm8, "net1", r=rc)
        C("net1", Von, c=cc)
        R(Vm7, "net2", r=rc)
        C("net2", Vop, c=cc)

        # MOS
        (
            Vdd()
            >> (M14 := PMOS(l=lp, w=wp, m=5*multi_p)) % Vm2
            >> Vm6
            >> (M17 := PMOS(l=lp1, w=wp1, m=1*multi_p1)) % Vcmfb
            >> (M18 := PMOS(l=lp2, w=wp2, m=1*multi_p2)) % VB2
            >> Vm5
        )
        Vdd() >> (M13 := PMOS(l=lp, w=wp, m=8*multi_p)) % Vm2 >> Vtail
        Vm2 >> (MB9 := PMOS(l=lp2, w=wp2, m=1*multi_p2)) % VB2 >> VB2

        # layout
        bias = col(
            row(cmp1b.pguard(), cmp1.pguard(), spacing=2000),
            layout_mos(MB9).pguard(),
            row(cmn1.nguard(),cmn2b.nguard().move_by(2000,0)),
            spacing=2000,
        )

        cmfb = col(
            layout_mos(M14).pguard(),
            row(layout_mos(M17).pguard(), dp7.pguard(), spacing=2000),
            layout_mos(M18).pguard(),
            cmn2.nguard(),
            spacing=2000,
        )

        stage1_2 = col(
            row(layout_mos(M13).pguard(), dp5.pguard(), spacing=2000),
            dp1.pguard(),
            dp2.pguard(),
            dp3.nguard(),
            row(dp4.nguard(), dp6.nguard(), spacing=2000),
            spacing=2000,
            align=LEFT
            )
        
        #stage2 = col(dp5.pguard(), dp6.nguard(), spacing=2000, align=CENTER)

        row(bias, cmfb, stage1_2,  spacing=2000, align=TOP)

        place()
        route()

    def wrapper(self, VINP, VINN, VCM, VON: Output, VOP: Output):
        IDC = Node()
        # self.build(VINP, VINN, IDC, VCM, VOP, VON)
        (VINP, VINN) >> self.build() % (IDC, VCM) >> (VOP, VON)
        # V(VCM, Vss(), dc=0.6)
        I(Vdd(), IDC, dc=25e-6)
        C(VOP, Vss(), c=3e-12)
        C(VON, Vss(), c=3e-12)



"""
说明
运行生成器 test_sim
运行优化器 test_opt

introduction
run test_sim to generate a fdamp and test it
run test_opt to genetate a fdamp and optimize it simultaneously as well as testing it
"""
def test_layout():
    @testbench
    def fdamp_layout():
        lay(FDAmp(vdd=1.2, Av=1e4, GBW=100e6, SR=30, PM_deg=60, power=2e-3, Cl=3e-12))



def test_sim():
    # corners = ["tt_lib", "ss_lib", "ff_lib"]
    # corners = ["tt_lib"]
    # for i in corners:
    #     with get_circuit().setenv(corners=[i]):
    #         raw = FDAmp(vdd=1.2, Av=1e4, GBW=100e6, SR=30, PM_deg=63, power=2e-3, Cl=3e-12).test()
    raw = FDAmp(vdd=1.2, Av=1e4, GBW=100e6, SR=30,
                PM_deg=60, power=2e-3, Cl=3e-12).test()


def test_opt():
    #    corners = ["ff_lib"]
    #    for i in corners:
    #        with get_circuit().setenv(corners=[i]):
    from frontend.tester import CombinedTester
    from taco.ea import DE

    ota = FDAmp(vdd=1.2, Av=1e4, GBW=100e6, SR=30,
                PM_deg=60, power=2e-3, Cl=3e-12)
    ota_tester = ota.test()
    tester = CombinedTester(
        [
            ota_tester.gainTester,
            ota_tester.noiseTester,
            ota_tester.DCTester,
            #               ota_tester.psrrTester,
            #               ota_tester.cmrrTester,
            ota_tester.srTester,
            #               ota_tester.THDTester,
            #                ota_tester.ICMRTester,
        ]
    )

    def problem(tester):
        from taco.problem import Problem

        gain = tester["gain"]
        bw = tester["bw"]
        gbw = gain * bw
        pm = tester["pm"]
#           ugb = tester["ugb"]
        power = tester["power"]
        gainDB = gain.db()
#           cmrr = abs((gain / tester["cmrr"]).db())
#           psrr = abs((gain / tester["psrr"]).db())
        srp, srn = tester["SRp"], tester["SRn"]
        thermal, f1 = tester["thermal"], tester["f1"]
#           thd = tester["thd"]
#          ICMRmin_OPE = tester["ICMRmin_OPE"]
#         ICMRmax_OPE = tester["ICMRmax_OPE"]
        debug(
            "\n"
            + bcolors.render(
                ", ".join(
                    [
                        f"power = {power*1e3:.2e}mW",
                        f"gain = {gainDB:.2f}dB",
                        f"gbw = {gbw/1e6:.2f}MHz",
                        #                        f"ugb = {ugb/1e6:.4f}MHz",
                        #                        f"bw = {bw/1e3:.2f}KHz",
                        f"pm = {pm:.2f}",
                        #                         f"cmrr={cmrr:2f}dB",
                        #                        f"psrr={psrr:.2f}dB",
                        f"srp={srp*1e-6:2f}V/us",
                        f"srn={srn*1e-6:2f}V/us",
                        f"thermal noise = {thermal*1e9:.2f}nV/Sqrt(Hz)",
                        f"f1 noise = {f1*1e9:.2f}nV/Sqrt(Hz)",
                        #                         f"Total harmonic distortion(THD) = {thd:.2e} %",
                        #                         f"ICMRmin_OPE = {ICMRmin_OPE:.2f} V",
                        #                         f"ICMRmax_OPE = {ICMRmax_OPE:.2f} V",
                    ]
                ),
                color="yellow",
            )
        )
        return Problem(
            objective=(min(srp, -srn)).maximize().rename("SR"),
            constraint=(
                power.lt(2e-3).rename("Power"),
                gbw.gt(100e6).rename("GBW"),
                #                   ugb.gt(80e6).rename("UGB"),
                gainDB.gt(80).rename("Av"),
                #                    bw.gt(3e3).rename("BW"),
                pm.gt(60).rename("PM"),
                pm.lt(90).rename("PM"),
                #                    cmrr.gt(100).rename("CMRR"),
                #                    psrr.gt(100).rename("PSRR"),
                srp.gt(30e6).rename("SRp"),
                srn.lt(-30e6).rename("SRn"),
                #                    f1.lt(100e-9).rename("f1 Noise"),
                thermal.lt(50e-9).rename("thermal Noise"),
                #                   thd.lt(0.1).rename("thd"),
                #                   ICMRmin_OPE.lt(0.4).rename("ICMRmin"),
                #                   ICMRmax_OPE.gt(0.8).rename("ICMRmax"),
            ),
        )

    configs = {
        "pop_size": 100,
        "maxgen": 100,
        "mutOper_F": 0.8,
        "recOper_XOVR": 0.5
    }
    opt = tester.optimize(
        problem=problem,
        optimizer=DE,
        **configs,
    )
#       from benchmarks.utils import local_perturbation
#        prophetVars = local_perturbation(configs["pop_size"])
    res = opt.solve(prophetVars=None)
    if res["success"]:
        ota.report(res=res, is_opt=True)

# pytest -s -vv test_fd_amp3.py::test_opt
